using System;
using System.Collections.Generic;
using System.Linq;
using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Volo.Abp.DependencyInjection;

namespace CKY.AgentPlatform.Domain.McpTools
{
    /// <summary>
    /// MCP JSON-RPC消息类型
    /// </summary>
    public enum McpMessageType
    {
        Request = 0,
        Response = 1,
        Notification = 2
    }

    /// <summary>
    /// MCP JSON-RPC消息基类
    /// </summary>
    public abstract class McpMessage
    {
        /// <summary>
        /// JSON-RPC版本
        /// </summary>
        public string JsonRpc { get; set; } = "2.0";

        /// <summary>
        /// 消息类型
        /// </summary>
        public McpMessageType Type { get; set; }

        /// <summary>
        /// 请求ID（用于Request/Response类型）
        /// </summary>
        public string Id { get; set; }
    }

    /// <summary>
    /// MCP请求消息
    /// </summary>
    public class McpRequest : McpMessage
    {
        /// <summary>
        /// 方法名
        /// </summary>
        public string Method { get; set; }

        /// <summary>
        /// 参数
        /// </summary>
        public object Params { get; set; }

        public McpRequest()
        {
            Type = McpMessageType.Request;
        }
    }

    /// <summary>
    /// MCP响应消息
    /// </summary>
    public class McpResponse : McpMessage
    {
        /// <summary>
        /// 结果数据
        /// </summary>
        public object Result { get; set; }

        /// <summary>
        /// 错误信息
        /// </summary>
        public McpError Error { get; set; }

        public McpResponse()
        {
            Type = McpMessageType.Response;
        }
    }

    /// <summary>
    /// MCP错误信息
    /// </summary>
    public class McpError
    {
        /// <summary>
        /// 错误代码
        /// </summary>
        public int Code { get; set; }

        /// <summary>
        /// 错误消息
        /// </summary>
        public string Message { get; set; }

        /// <summary>
        /// 错误详情
        /// </summary>
        public object Data { get; set; }
    }

    /// <summary>
    /// MCP通知消息
    /// </summary>
    public class McpNotification : McpMessage
    {
        /// <summary>
        /// 方法名
        /// </summary>
        public string Method { get; set; }

        /// <summary>
        /// 参数
        /// </summary>
        public object Params { get; set; }

        public McpNotification()
        {
            Type = McpMessageType.Notification;
        }
    }

    /// <summary>
    /// MCP工具定义
    /// </summary>
    public class McpToolDefinition
    {
        /// <summary>
        /// 工具名称
        /// </summary>
        public string Name { get; set; }

        /// <summary>
        /// 工具描述
        /// </summary>
        public string Description { get; set; }

        /// <summary>
        /// 输入模式
        /// </summary>
        public object InputSchema { get; set; }

        /// <summary>
        /// 工具配置
        /// </summary>
        public Dictionary<string, object> Configuration { get; set; }
    }

    /// <summary>
    /// MCP资源定义
    /// </summary>
    public class McpResourceDefinition
    {
        /// <summary>
        /// 资源URI
        /// </summary>
        public string Uri { get; set; }

        /// <summary>
        /// 资源名称
        /// </summary>
        public string Name { get; set; }

        /// <summary>
        /// 资源描述
        /// </summary>
        public string Description { get; set; }

        /// <summary>
        /// MIME类型
        /// </summary>
        public string MimeType { get; set; }
    }

    /// <summary>
    /// MCP客户端配置
    /// </summary>
    public class McpClientConfig
    {
        /// <summary>
        /// 服务器URL
        /// </summary>
        public string ServerUrl { get; set; }

        /// <summary>
        /// 连接超时时间（毫秒）
        /// </summary>
        public int ConnectTimeoutMs { get; set; } = 30000;

        /// <summary>
        /// 心跳间隔（秒）
        /// </summary>
        public int HeartbeatIntervalSeconds { get; set; } = 30;

        /// <summary>
        /// 最大重试次数
        /// </summary>
        public int MaxRetryCount { get; set; } = 3;

        /// <summary>
        /// 是否启用调试日志
        /// </summary>
        public bool EnableDebugLog { get; set; } = false;

        /// <summary>
        /// 自定义请求头
        /// </summary>
        public Dictionary<string, string> Headers { get; set; } = new Dictionary<string, string>();
    }

    /// <summary>
    /// MCP客户端事件委托
    /// </summary>
    /// <param name="sender">发送者</param>
    /// <param name="message">消息</param>
    public delegate void McpMessageEventHandler(object sender, McpMessage message);

    /// <summary>
    /// MCP连接状态变更事件委托
    /// </summary>
    /// <param name="sender">发送者</param>
    /// <param name="isConnected">是否连接</param>
    public delegate void McpConnectionStateEventHandler(object sender, bool isConnected);

    /// <summary>
    /// MCP工具列表变更事件委托
    /// </summary>
    /// <param name="sender">发送者</param>
    /// <param name="tools">工具列表</param>
    public delegate void McpToolsEventHandler(object sender, List<McpToolDefinition> tools);

    /// <summary>
    /// MCP资源列表变更事件委托
    /// </summary>
    /// <param name="sender">发送者</param>
    /// <param name="resources">资源列表</param>
    public delegate void McpResourcesEventHandler(object sender, List<McpResourceDefinition> resources);

    /// <summary>
    /// MCP客户端 - 实现与MCP服务器的通信
    /// </summary>
    public class McpClient : ITransientDependency
    {
        private readonly McpClientConfig _config;
        private ClientWebSocket _webSocket;
        private CancellationTokenSource _cancellationTokenSource;
        private readonly Dictionary<string, TaskCompletionSource<McpResponse>> _pendingRequests = new Dictionary<string, TaskCompletionSource<McpResponse>>();
        private readonly List<McpToolDefinition> _availableTools = new List<McpToolDefinition>();
        private readonly List<McpResourceDefinition> _availableResources = new List<McpResourceDefinition>();
        private bool _isConnected = false;
        private int _requestIdCounter = 1;

        // 事件
        public event McpMessageEventHandler OnMessageReceived;
        public event McpConnectionStateEventHandler OnConnectionStateChanged;
        public event McpToolsEventHandler OnToolsUpdated;
        public event McpResourcesEventHandler OnResourcesUpdated;
        public event EventHandler<Exception> OnError;

        public McpClient(McpClientConfig config)
        {
            _config = config;
        }

        /// <summary>
        /// 连接到MCP服务器
        /// </summary>
        public async Task ConnectAsync()
        {
            try
            {
                _cancellationTokenSource = new CancellationTokenSource();
                _webSocket = new ClientWebSocket();

                // 添加自定义请求头
                foreach (var header in _config.Headers)
                {
                    _webSocket.Options.SetRequestHeader(header.Key, header.Value);
                }

                // 连接到服务器
                await _webSocket.ConnectAsync(new Uri(_config.ServerUrl), _cancellationTokenSource.Token);
                _isConnected = true;
                OnConnectionStateChanged?.Invoke(this, true);

                if (_config.EnableDebugLog)
                {
                    Console.WriteLine($"Connected to MCP server: {_config.ServerUrl}");
                }

                // 启动消息接收循环
                _ = Task.Run(() => ReceiveMessagesAsync(_cancellationTokenSource.Token));

                // 启动心跳循环
                _ = Task.Run(() => HeartbeatLoopAsync(_cancellationTokenSource.Token));

                // 初始化：获取工具和资源列表
                await InitializeAsync();
            }
            catch (Exception ex)
            {
                _isConnected = false;
                OnError?.Invoke(this, ex);
                throw new McpConnectionException("Failed to connect to MCP server", ex);
            }
        }

        /// <summary>
        /// 断开连接
        /// </summary>
        public async Task DisconnectAsync()
        {
            try
            {
                _cancellationTokenSource?.Cancel();
                
                if (_webSocket?.State == WebSocketState.Open)
                {
                    await _webSocket.CloseAsync(
                        WebSocketCloseStatus.NormalClosure,
                        "Disconnecting",
                        CancellationToken.None);
                }

                _isConnected = false;
                OnConnectionStateChanged?.Invoke(this, false);

                // 清理待处理的请求
                foreach (var request in _pendingRequests.Values)
                {
                    request.TrySetCanceled();
                }
                _pendingRequests.Clear();
            }
            catch (Exception ex)
            {
                OnError?.Invoke(this, ex);
            }
            finally
            {
                _webSocket?.Dispose();
                _webSocket = null;
            }
        }

        /// <summary>
        /// 发送请求到MCP服务器
        /// </summary>
        /// <param name="method">方法名</param>
        /// <param name="parameters">参数</param>
        /// <returns>响应结果</returns>
        public async Task<McpResponse> SendRequestAsync(string method, object parameters = null)
        {
            if (!_isConnected)
            {
                throw new McpConnectionException("Not connected to MCP server");
            }

            var requestId = (_requestIdCounter++).ToString();
            var request = new McpRequest
            {
                Id = requestId,
                Method = method,
                Params = parameters
            };

            var tcs = new TaskCompletionSource<McpResponse>();
            _pendingRequests[requestId] = tcs;

            try
            {
                await SendMessageAsync(request);
                return await tcs.Task;
            }
            finally
            {
                _pendingRequests.Remove(requestId);
            }
        }

        /// <summary>
        /// 发送通知到MCP服务器
        /// </summary>
        /// <param name="method">方法名</param>
        /// <param name="parameters">参数</param>
        public async Task SendNotificationAsync(string method, object parameters = null)
        {
            if (!_isConnected)
            {
                throw new McpConnectionException("Not connected to MCP server");
            }

            var notification = new McpNotification
            {
                Method = method,
                Params = parameters
            };

            await SendMessageAsync(notification);
        }

        /// <summary>
        /// 调用工具
        /// </summary>
        /// <param name="toolName">工具名称</param>
        /// <param name="parameters">工具参数</param>
        /// <returns>工具执行结果</returns>
        public async Task<object> CallToolAsync(string toolName, object parameters = null)
        {
            var response = await SendRequestAsync("tools/call", new
            {
                name = toolName,
                arguments = parameters
            });

            if (response.Error != null)
            {
                throw new McpToolException($"Tool execution failed: {response.Error.Message}", response.Error);
            }

            return response.Result;
        }

        /// <summary>
        /// 读取资源
        /// </summary>
        /// <param name="resourceUri">资源URI</param>
        /// <returns>资源内容</returns>
        public async Task<object> ReadResourceAsync(string resourceUri)
        {
            var response = await SendRequestAsync("resources/read", new
            {
                uri = resourceUri
            });

            if (response.Error != null)
            {
                throw new McpResourceException($"Resource read failed: {response.Error.Message}", response.Error);
            }

            return response.Result;
        }

        /// <summary>
        /// 获取可用工具列表
        /// </summary>
        /// <returns>工具列表</returns>
        public List<McpToolDefinition> GetAvailableTools()
        {
            return _availableTools.ToList();
        }

        /// <summary>
        /// 获取可用资源列表
        /// </summary>
        /// <returns>资源列表</returns>
        public List<McpResourceDefinition> GetAvailableResources()
        {
            return _availableResources.ToList();
        }

        /// <summary>
        /// 检查是否已连接
        /// </summary>
        public bool IsConnected() => _isConnected;

        private async Task InitializeAsync()
        {
            try
            {
                // 获取工具列表
                var toolsResponse = await SendRequestAsync("tools/list");
                if (toolsResponse.Result is JsonElement toolsElement)
                {
                    _availableTools.Clear();
                    var tools = JsonSerializer.Deserialize<List<McpToolDefinition>>(toolsElement.GetRawText());
                    if (tools != null)
                    {
                        _availableTools.AddRange(tools);
                        OnToolsUpdated?.Invoke(this, _availableTools);
                    }
                }

                // 获取资源列表
                var resourcesResponse = await SendRequestAsync("resources/list");
                if (resourcesResponse.Result is JsonElement resourcesElement)
                {
                    _availableResources.Clear();
                    var resources = JsonSerializer.Deserialize<List<McpResourceDefinition>>(resourcesElement.GetRawText());
                    if (resources != null)
                    {
                        _availableResources.AddRange(resources);
                        OnResourcesUpdated?.Invoke(this, _availableResources);
                    }
                }
            }
            catch (Exception ex)
            {
                OnError?.Invoke(this, new McpInitializationException("Failed to initialize MCP client", ex));
            }
        }

        private async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
        {
            var buffer = new byte[8192];
            var messageBuffer = new List<byte>();

            try
            {
                while (_webSocket?.State == WebSocketState.Open && !cancellationToken.IsCancellationRequested)
                {
                    WebSocketReceiveResult result;
                    do
                    {
                        result = await _webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), cancellationToken);
                        if (result.Count > 0)
                        {
                            messageBuffer.AddRange(buffer.Take(result.Count));
                        }
                    } while (!result.EndOfMessage);

                    if (messageBuffer.Count > 0)
                    {
                        var messageJson = Encoding.UTF8.GetString(messageBuffer.ToArray());
                        var message = ParseMessage(messageJson);

                        if (_config.EnableDebugLog)
                        {
                            Console.WriteLine($"Received MCP message: {messageJson}");
                        }

                        await ProcessMessageAsync(message);
                        messageBuffer.Clear();
                    }
                }
            }
            catch (OperationCanceledException)
            {
                // 正常关闭
            }
            catch (Exception ex)
            {
                OnError?.Invoke(this, ex);
            }
            finally
            {
                _isConnected = false;
                OnConnectionStateChanged?.Invoke(this, false);
            }
        }

        private async Task HeartbeatLoopAsync(CancellationToken cancellationToken)
        {
            while (!cancellationToken.IsCancellationRequested && _isConnected)
            {
                try
                {
                    await Task.Delay(TimeSpan.FromSeconds(_config.HeartbeatIntervalSeconds), cancellationToken);
                    
                    if (_isConnected)
                    {
                        await SendNotificationAsync("ping");
                    }
                }
                catch (OperationCanceledException)
                {
                    // 正常关闭
                }
                catch (Exception ex)
                {
                    OnError?.Invoke(this, ex);
                }
            }
        }

        private async Task SendMessageAsync(McpMessage message)
        {
            if (_webSocket?.State != WebSocketState.Open)
            {
                throw new McpConnectionException("WebSocket connection is not open");
            }

            var messageJson = JsonSerializer.Serialize(message, new JsonSerializerOptions
            {
                PropertyNamingPolicy = JsonNamingPolicy.CamelCase
            });

            if (_config.EnableDebugLog)
            {
                Console.WriteLine($"Sending MCP message: {messageJson}");
            }

            var buffer = Encoding.UTF8.GetBytes(messageJson);
            await _webSocket.SendAsync(
                new ArraySegment<byte>(buffer),
                WebSocketMessageType.Text,
                true,
                CancellationToken.None);
        }

        private McpMessage ParseMessage(string messageJson)
        {
            var jsonElement = JsonSerializer.Deserialize<JsonElement>(messageJson);

            if (jsonElement.TryGetProperty("method", out _))
            {
                if (jsonElement.TryGetProperty("id", out _))
                {
                    return JsonSerializer.Deserialize<McpRequest>(messageJson);
                }
                else
                {
                    return JsonSerializer.Deserialize<McpNotification>(messageJson);
                }
            }
            else
            {
                return JsonSerializer.Deserialize<McpResponse>(messageJson);
            }
        }

        private async Task ProcessMessageAsync(McpMessage message)
        {
            try
            {
                OnMessageReceived?.Invoke(this, message);

                if (message is McpResponse response)
                {
                    if (_pendingRequests.TryGetValue(response.Id, out var tcs))
                    {
                        tcs.TrySetResult(response);
                    }
                }
                else if (message is McpNotification notification)
                {
                    await ProcessNotificationAsync(notification);
                }
            }
            catch (Exception ex)
            {
                OnError?.Invoke(this, ex);
            }
        }

        private async Task ProcessNotificationAsync(McpNotification notification)
        {
            switch (notification.Method)
            {
                case "tools/list_changed":
                    await InitializeAsync(); // 重新获取工具列表
                    break;
                    
                case "resources/list_changed":
                    await InitializeAsync(); // 重新获取资源列表
                    break;
            }
        }
    }

    /// <summary>
    /// MCP连接异常
    /// </summary>
    public class McpConnectionException : Exception
    {
        public McpConnectionException(string message) : base(message) { }
        public McpConnectionException(string message, Exception innerException) : base(message, innerException) { }
    }

    /// <summary>
    /// MCP工具异常
    /// </summary>
    public class McpToolException : Exception
    {
        public McpToolException(string message) : base(message) { }
        public McpToolException(string message, McpError error) : base(message) { }
    }

    /// <summary>
    /// MCP资源异常
    /// </summary>
    public class McpResourceException : Exception
    {
        public McpResourceException(string message) : base(message) { }
        public McpResourceException(string message, McpError error) : base(message) { }
    }

    /// <summary>
    /// MCP初始化异常
    /// </summary>
    public class McpInitializationException : Exception
    {
        public McpInitializationException(string message) : base(message) { }
        public McpInitializationException(string message, Exception innerException) : base(message, innerException) { }
    }
}